from itertools import count
from os import path

import torch
from torch.autograd import Variable

import DQN
from Agent import Agent
from env import FlappyEnvironment
from  matplotlib import pyplot as plt
import numpy as np

from utilenn import tensor_image_to_numpy_image

env = FlappyEnvironment()

model = DQN.DQN()

if torch.cuda.is_available():
    model.cuda()

if path.exists('./dqn.net'):
    model.load_state_dict(torch.load('./dqn.net'))

env.reset()

env.step(1)
env.step(0)
env.step(1)

env.step(0)
env.step(1)

env.step(0)
env.step(1)

current_state = env.current_state

cnns = model.conv1

allout = cnns.forward(Variable(current_state))
cmap = 'gist_heat'

for i in range(4):

    if i > 0:
        layers = []
        for l in range(i):
            layers.append(cnns[l])

        net = torch.nn.Sequential(*layers)

        out = net(Variable(current_state))
        tensor = out.data
        _, channels, _, _, = tensor.size()

        channels = min(channels, 5)

        f, axarr = plt.subplots(1, channels)

        for c in range(channels):
            axarr[c].imshow(tensor[0][c].cpu().numpy(), cmap=cmap)

plt.show(block=True)
